import json
import tornado.web
import LoggerDefault
from litie.pipelines import UIEPipeline

# 实体识别
schema = ['时间'] 
# uie-base模型已上传至huggingface，可自动下载，其他模型只需提供模型名称将自动进行转换
uie = UIEPipeline("uie_lit_model_pytorch", schema=schema,device="cpu")

import numpy
class NumpyEncoder(json.JSONEncoder):  
    def default(self, obj):  
        if isinstance(obj, (numpy.int_, numpy.intc, numpy.intp, numpy.int8,  
            numpy.int16, numpy.int32, numpy.int64, numpy.uint8,  
            numpy.uint16, numpy.uint32, numpy.uint64)):  
            return int(obj)  
        elif isinstance(obj, (numpy.float_, numpy.float16, numpy.float32,numpy.float64)):  
            return float(obj)  
        elif isinstance(obj, (numpy.ndarray,)):  
            return obj.tolist()  
        return json.JSONEncoder.default(self, obj) 


class ClassifyRecognition(tornado.web.RequestHandler):
        async def post(self):
            #利用ｒｅｑｕｅｓｔ属性
            #取出客户端提交的ｊｓｏｎ字符串
            jsonbyte = self.request.body
            jsonstr = jsonbyte.decode('utf8')  #解码，二进制转为字符串
            jsonobj = json.loads(jsonstr)  #将字符串转为json对象
            texts = jsonobj.get('texts')#就可以用api取值
            schemas = jsonobj.get('schemas')#就可以用api取值
            uie.set_schema(schemas)
            maxlength = jsonobj.get('maxLength')#就可以用api取值
            prob = jsonobj.get('prob')#就可以用api取值
            result = uie(texts)
            res = {'result':{'code':"1", 'message':'success'},'body':result}
            self.write(json.dumps(res,cls=NumpyEncoder,ensure_ascii=False))

class Application(tornado.web.Application):
    _routes = [
        tornado.web.url(r'/extract', ClassifyRecognition)
    ]

    def __init__(self):
        super(Application, self).__init__(self._routes)

if __name__ == "__main__":
    application=Application()
    application.listen(12008)
    LoggerDefault.logging.info("start 12008")
    tornado.ioloop.IOLoop.instance().start() 